package cn.itcast.tags.tools

import cn.itcast.tags.utils.HdfsUtils
import org.apache.hadoop.conf.Configuration
import org.apache.hadoop.fs.Hdfs
import org.apache.spark.internal.Logging
import org.apache.spark.ml.Model
import org.apache.spark.ml.classification.DecisionTreeClassificationModel
import org.apache.spark.ml.clustering.{KMeans, KMeansModel}
import org.apache.spark.sql.DataFrame
import org.apache.spark.storage.StorageLevel

object MLModelTools extends Logging{

  def trainBestKMeansModel(dataframe: DataFrame, k: Int = 2): KMeansModel = {

    val maxIters:Array[Int] =Array(5,10,20)
    dataframe.persist(StorageLevel.MEMORY_AND_DISK)

    val models: Array[(Double, KMeansModel, Int)] = maxIters.map {
      maxIter => {
        val kMeans: KMeans = new KMeans()
          .setFeaturesCol("features")
          .setPredictionCol("predicition")
          .setK(4)
          .setMaxIter(maxIter)
          .setSeed(4869L)
        val model: KMeansModel = kMeans.fit(dataframe)
        val sse: Double = model.computeCost(dataframe)
        (sse, model, maxIter)
      }
    }
    dataframe.unpersist()
    models.foreach(println)
    val (_,bestModel,_)= models.minBy(_._1)
    bestModel
  }

  def  loadModel(dataframe:DataFrame, mlType:String, modelPath:String):Model[_]={
    val conf: Configuration = dataframe.sparkSession.sparkContext.hadoopConfiguration
    if(HdfsUtils.exists(conf,modelPath)){
      logWarning(s"loading model from ${modelPath}.........")
      mlType.toLowerCase() match {
        case "rfm" => KMeansModel.load(modelPath)
        case "rfe" => KMeansModel.load(modelPath)
        case "psm" => KMeansModel.load(modelPath)
        case "usg" => DecisionTreeClassificationModel.load(modelPath)

      }
    }else{
      logWarning(s"is training model .........")
      val bestModel:KMeansModel = mlType.toLowerCase() match {
        case "rfm" => trainBestKMeansModel(dataframe, k = 5)
        case "rfe" => trainBestKMeansModel(dataframe, k = 4)
        case "psm" => trainBestKMeansModel(dataframe, k = 5)
        case "usg" => null
      }

      logWarning(s"is saving model  for ${modelPath}.........")
      bestModel.save(modelPath)
      bestModel
    }

  }

}
